# 验证阶段剖面图绘制，云顶高度在上，剖面图在下
import numpy as np
import numpy.ma as ma
import netCDF4 as nc
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
import pandas as pd
from datetime import datetime
from scipy.spatial import cKDTree
from read_CloudSat import reader

# ------------------------ 云顶高度加载与预处理 ------------------------
def load_and_preprocess_cth(file_path):
    with nc.Dataset(file_path, 'r') as nf:
        cth = np.ma.getdata(nf.variables['cth'][:])
        lat = nf.variables['lat'][:]
        lon = nf.variables['lon'][:]

    cth = np.ma.masked_invalid(cth).filled(np.nan) / 1000.0
    lat = np.array(lat)
    lon = np.array(lon)

    # 处理 lat/lon 为2D网格
    if lat.ndim == 1 and lon.ndim == 1:
        lon, lat = np.meshgrid(lon, lat)

    # 替换 lat/lon 中的 NaN/Inf 为均值（避免 pcolormesh 报错）
    lat = np.where(np.isfinite(lat), lat, np.nanmean(lat))
    lon = np.where(np.isfinite(lon), lon, np.nanmean(lon))
    cth = np.where(np.isfinite(cth), cth, np.nan)

    return cth, lat, lon


# ------------------------ 绘制云顶高度分布图（顶部圆形地图） ------------------------
def draw_cloud_top_map(ax, cth, lat, lon):
    # 强制类型为 ndarray 且不含 mask
    lat = np.array(lat)
    lon = np.array(lon)
    cth = np.array(cth)

    # 避免非有限值报错
    cth = np.where(np.isfinite(cth), cth, np.nan)
    lat = np.where(np.isfinite(lat), lat, np.nanmean(lat))
    lon = np.where(np.isfinite(lon), lon, np.nanmean(lon))

    global_min = np.nanmin(cth)
    global_max = np.nanmax(cth)

    im = ax.pcolormesh(
        lon, lat, cth,
        transform=ccrs.PlateCarree(),
        cmap='jet',
        vmin=global_min,
        vmax=global_max
    )

    ax.set_global()
    ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5, edgecolor='black')
    ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.4, edgecolor='black')
    ax.set_title('Cloud Top Height (km)', fontsize=10)
    return im


# ------------------------ CloudSat轨迹绘制 ------------------------
def draw_track(ax, lon, lat, lat_range=None):
    mask = (lat >= lat_range[0]) & (lat <= lat_range[1]) if lat_range else np.ones_like(lat, dtype=bool)
    ax.plot(lon[mask], lat[mask], lw=2, color='r', transform=ccrs.Geodetic())
    if np.any(mask):
        ax.plot(lon[mask][0], lat[mask][0], 'ro', ms=3, transform=ccrs.PlateCarree())
        ax.text(lon[mask][0] + 5, lat[mask][0], 'start', color='r', fontsize='x-small', transform=ccrs.PlateCarree())

# ------------------------ 剖面图绘制 ------------------------
def draw_cross_section(ax, lon, hgt, data, cth_sorted, cbh_sorted, latitude_sorted, xlims=None):
    im = ax.pcolormesh(lon, hgt, data, cmap='jet', shading='nearest')
    ax.set_ylim(0, 20)
    ax.set_xlim(xlims if xlims else (lon.min(), lon.max()))
    ax.set_xlabel('Latitude', fontsize='x-small')
    ax.set_ylabel('Height [km]', fontsize='x-small')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('bottom', size='10%', pad=0.42)
    cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
    cbar.ax.tick_params(labelsize='x-small')
    cbar.set_label('dBZe', fontsize='x-small')
    ax.plot(latitude_sorted, cth_sorted, color='black', linewidth=0.8, label='CTH')
    ax.plot(latitude_sorted, cbh_sorted, color='red', linewidth=0.8, label='CBH')
    ax.legend(fontsize='x-small')

# ------------------------ 地形绘制 ------------------------
def draw_elevation(ax, lon, elv):
    ax.fill_between(lon, elv, color='gray')

# ------------------------ 主程序 ------------------------
if __name__ == '__main__':
    # 文件路径
    cloudsat_path = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020139071017_74881_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf' #5.18
    # cloudsat_path = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020159081830_75174_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf' #6.7
    # cloudsat_path = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020142040823_74923_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'  #5.21
    pred_path = '/mnt/datastore/liudddata/result/20200506_droupout/2020051808_predicted_2d_mc.nc'
    # pred_path = '/mnt/datastore/liudddata/result/20200506_droupout/2020052105_predicted_2d_mc.nc'
    # pred_path = '/mnt/datastore/liudddata/result/20200506_droupout/2020060709_predicted_2d_mc.nc'
    coord_file = '/home/liudd/data_preprocessing/FY4A_coordinates.nc'

    # 读取 CloudSat 数据
    f = reader(cloudsat_path)
    lon, lat, elv = f.read_geo()
    data = f.read_sds('Radar_Reflectivity')
    height = f.read_sds('Height')
    time = f.read_time(datetime=True)
    f.close()

    # start_dt = datetime(2020, 5, 21, 5, 0)
    # end_dt = datetime(2020, 5, 21, 5, 15)
    start_dt = datetime(2020, 5, 18, 8, 0)
    end_dt = datetime(2020, 5, 18, 8, 15)
    # start_dt = datetime(2020, 6, 7, 9, 0)
    # end_dt = datetime(2020, 6, 7, 9, 15)
    time_series = pd.to_datetime(time)
    start_idx, end_idx = np.searchsorted(time_series, [start_dt, end_dt])
    rad_refl = np.transpose(data[start_idx:end_idx])
    rad_refl[rad_refl <= -20] = np.nan
    hgt = np.mean(height[start_idx:end_idx], axis=0) / 1000
    lat_t = lat[start_idx:end_idx]
    lon_t = lon[start_idx:end_idx]
    elv = elv / 1000

    # 匹配云顶云底高度
    pred_nc = nc.Dataset(pred_path, 'r')
    cth_data = pred_nc.variables['cth'][:]
    cbh_data = pred_nc.variables['predicted_mean'][:]

    coord_nc = nc.Dataset(coord_file, 'r')
    lat_fy = coord_nc.variables['lat'][:, :].T
    lon_fy = coord_nc.variables['lon'][:, :].T
    coord_nc.close()

    lat_fy[np.isnan(lat_fy)] = -90
    lon_fy[np.isnan(lon_fy)] = 360
    fy_coords = np.deg2rad(np.vstack([lat_fy.ravel(), lon_fy.ravel()]).T)
    fy_cartesian = np.array([
        np.cos(fy_coords[:, 0]) * np.cos(fy_coords[:, 1]),
        np.cos(fy_coords[:, 0]) * np.sin(fy_coords[:, 1]),
        np.sin(fy_coords[:, 0])
    ]).T
    tree = cKDTree(fy_cartesian[np.all(np.isfinite(fy_cartesian), axis=1)])
    radius = 2 / 6371

    latitudes, longitudes, cth_values, cbh_values = [], [], [], []
    for la, lo in zip(lat_t, lon_t):
        p = np.deg2rad([la, lo])
        p_cartesian = [np.cos(p[0]) * np.cos(p[1]), np.cos(p[0]) * np.sin(p[1]), np.sin(p[0])]
        indices = tree.query_ball_point(p_cartesian, r=radius)
        if indices:
            for idx in indices:
                lat_idx, lon_idx = np.unravel_index(idx, lat_fy.shape)
                latitudes.append(lat_fy[lat_idx, lon_idx])
                longitudes.append(lon_fy[lat_idx, lon_idx])
                cth_values.append(cth_data[lat_idx, lon_idx] * 0.001)
                cbh_values.append(cbh_data[lat_idx, lon_idx] * 0.001)
        else:
            latitudes.append(np.nan); longitudes.append(np.nan)
            cth_values.append(np.nan); cbh_values.append(np.nan)

    df = pd.DataFrame({'Latitude': latitudes, 'Longitude': longitudes, 'CTH': cth_values, 'CBH': cbh_values})
    df.dropna(inplace=True)
    df_sorted = df.sort_values(by='Latitude')
    latitude_sorted = df_sorted['Latitude']
    cth_sorted = df_sorted['CTH']
    cbh_sorted = df_sorted['CBH']

    # 读取云图网格数据
    cth_map, lat_map, lon_map = load_and_preprocess_cth(pred_path)

    # 图像绘制
    fig = plt.figure(figsize=(5, 7), dpi=300)
    ax1 = fig.add_axes([0.15, 0.55, 0.7, 0.4], projection=ccrs.Orthographic(central_longitude=105, central_latitude=0))
    ax2 = fig.add_axes([0.15, 0.1, 0.7, 0.35]) #[left, bottom, width, height]

    # 绘图
    im = draw_cloud_top_map(ax1, cth_map, lat_map, lon_map)
    # draw_track(ax1, lon_t, lat_t, lat_range=(15, 33)) #5.21
    # draw_track(ax1, lon_t, lat_t, lat_range=(-28, -18.5)) #6.7
    draw_track(ax1, lon_t, lat_t, lat_range=(4, 18)) #5.18
    # draw_cross_section(ax2, lat_t, hgt, rad_refl, cth_sorted, cbh_sorted, latitude_sorted, xlims=(15, 33))
    # draw_cross_section(ax2, lat_t, hgt, rad_refl, cth_sorted, cbh_sorted, latitude_sorted, xlims=(-28, -18.5))
    draw_cross_section(ax2, lat_t, hgt, rad_refl, cth_sorted, cbh_sorted, latitude_sorted, xlims=(4, 18))
    draw_elevation(ax2, lat_t, elv[start_idx:end_idx])

    # 添加颜色条
    cbar = plt.colorbar(im, ax=ax1, orientation='horizontal', pad=0.05, shrink=0.8)
    cbar.set_label('Cloud Top Height (km)', fontsize=8)

    # 标题
    ax2.set_title('Radar Reflectivity & FY4A Cloud Heights', fontsize='small')
    ax1.set_title('Cloud Top Height & CloudSat Track', fontsize='small')

    # 保存和展示
    # fig.savefig('/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020052105_combined_figure.png', bbox_inches='tight')
    fig.savefig('/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020051809_combined_figure.png', bbox_inches='tight')
    # fig.savefig('/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/20200060709_combined_figure.png', bbox_inches='tight')
    plt.show()
